from kafka import KafkaProducer
from kafka.errors import KafkaError
import json
import time
import argparse
import logging
from typing import Dict, Any

# 配置日志格式
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class SecureKafkaProducer:
    def __init__(self, config: Dict[str, Any]):
        """
        初始化鉴权Kafka生产者
        :param config: 配置字典，包含以下字段：
            - bootstrap_servers: Kafka地址列表
            - topic: 目标topic
            - sasl_mechanism: 鉴权机制 (e.g. 'PLAIN')
            - sasl_plain_username: 用户名
            - sasl_plain_password: 密码
            - security_protocol: 安全协议 (e.g. 'SASL_SSL')
            - ssl_cafile: CA证书路径 (可选)
        """
        self.producer = KafkaProducer(
            bootstrap_servers=config['bootstrap_servers'],
            sasl_mechanism=config.get('sasl_mechanism', 'PLAIN'),
            sasl_plain_username=config.get('sasl_plain_username'),
            sasl_plain_password=config.get('sasl_plain_password'),
            security_protocol=config.get('security_protocol', 'SASL_PLAINTEXT'
                                                              ''),
            ssl_cafile=config.get('ssl_cafile'),
            value_serializer=lambda v: json.dumps(v).encode('utf-8'),
            acks='all',  # 确保消息可靠发送
            retries=3  # 失败重试
        )
        self.topic = config['topic']
        logger.info(f"Producer initialized for topic: {self.topic}")

    def send_batch(self, message_template: Dict[str, Any], batch_size: int) -> int:
        """
        发送一批消息
        :param message_template: 消息模板字典
        :param batch_size: 本批次消息数量
        :return: 成功发送的消息数
        """
        success = 0
        for i in range(batch_size):
            try:
                # 为每条消息添加唯一标识
                message = message_template.copy()
                message['msg_id'] = f"{int(time.time() * 1000)}_{i}"
                message['batch_seq'] = i + 1

                future = self.producer.send(
                    self.topic,
                    value=message
                )
                # 非阻塞获取发送结果（可选）
                # _ = future.get(timeout=10)
                success += 1
            except KafkaError as e:
                logger.error(f"Message {i} failed: {str(e)}")
        self.producer.flush()  # 确保本批次所有消息完成发送
        return success


def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='Secure Kafka Producer')
    parser.add_argument('--servers', required=True, help='Kafka bootstrap servers (comma separated)')
    parser.add_argument('--topic', required=True, help='Target topic name')
    parser.add_argument('--username', required=True, help='SASL username')
    parser.add_argument('--password', required=True, help='SASL password')
    parser.add_argument('--batch-size', type=int, default=10, help='Messages per batch')
    parser.add_argument('--iterations', type=int, default=5, help='Number of batches to send')
    parser.add_argument('--interval', type=float, default=1.0, help='Seconds between batches')
    parser.add_argument('--ca-file', help='Path to CA certificate file')
    args = parser.parse_args()

    # 构建Kafka配置
    kafka_config = {
        'bootstrap_servers': args.servers.split(','),
        'topic': args.topic,
        'sasl_plain_username': args.username,
        'sasl_plain_password': args.password,
        'ssl_cafile': args.ca_file
    }

    # 消息模板（可根据需要修改）
    message_template = {"metadata": {
                        "source": "python-producer",
                        "sequence": 0},
                        "timepoint.fmt.kafkaTimestampType": "CreateTime",
                        "field.in_param": [{"name": "ROOT.BODY.BUSI_INFO.COMMON_INFO.DEST_PORT", "value": "31000"},
                                           {"name": "ROOT.BODY.BUSI_INFO.COMMON_INFO.APP_NAME",
                                            "value": "windows_IE11"},
                                           {"name": "ROOT.BODY.BUSI_INFO.DESK_ID", "value": "null"},
                                           {"name": "ROOT.BODY.BUSI_INFO.COMMON_INFO.DEST_IP", "value": "10.243.39.18"},
                                           {"name": "ROOT.BODY.BUSI_INFO.COMMON_INFO.FROM_SYS", "value": ""},
                                           {"name": "ROOT.BODY.BUSI_INFO.OPR_INFO.CHANNEL_TYPE", "value": "431"},
                                           {"name": "ROOT.BODY.BUSI_INFO.OPR_INFO.RE_WAY", "value": ""},
                                           {"name": "ROOT.BODY.BUSI_INFO.COMMON_INFO.SESSION_ID",
                                            "value": "n466B175ABB796981712A82058B926575-1"},
                                           {"name": "ROOT.BODY.BUSI_INFO.OPR_INFO.OP_CODE", "value": ""},
                                           {"name": "ROOT.BODY.BUSI_INFO.OPR_INFO.REL_OP_CODE", "value": ""},
                                           {"name": "ROOT.BODY.BUSI_INFO.OPR_INFO.LOGIN_PWD", "value": ""},
                                           {"name": "ROOT.BODY.BUSI_INFO.OPR_INFO.AUTHEN_CODE", "value": ""},
                                           {"name": "ROOT.BODY.BUSI_INFO.COMMON_INFO.PROVINCE_GROUP", "value": "10017"},
                                           {"name": "ROOT.BODY.BUSI_INFO.OPR_INFO.REGION_ID", "value": "15"},
                                           {"name": "ROOT.BODY.BUSI_INFO.OPR_INFO.IP_ADDRESS",
                                            "value": "82.221.145.202"},
                                           {"name": "ROOT.BODY.BUSI_INFO.OPR_INFO.CONTACT_ID", "value": ""},
                                           {"name": "ROOT.BODY.BUSI_INFO.COMMON_INFO.CALL_ID",
                                            "value": "2c6b99bb-2668-46e8-bce5-e961eaf6f469"},
                                           {"name": "ROOT.BODY.BUSI_INFO.OPR_INFO.AUTHEN_NAME", "value": ""},
                                           {"name": "ROOT.BODY.BUSI_INFO.OPR_INFO.CHANNEL_SOURCE", "value": ""},
                                           {"name": "ROOT.BODY.BUSI_INFO.OPR_INFO.LOGIN_NO", "value": "E3ZZZZ063"},
                                           {"name": "ROOT.BODY.BUSI_INFO.OPR_INFO.GROUP_ID", "value": "115010975"},
                                           {"name": "ROOT.BODY.BUSI_INFO.COMMON_INFO.SER_NAME",
                                            "value": "tangmaolong01"},
                                           {"name": "ROOT.BODY.BUSI_INFO.COMMON_INFO.APP_ID", "value": ""},
                                           {"name": "ROOT.BODY.BUSI_INFO.COMMON_INFO.BACK2", "value": ""},
                                           {"name": "ROOT.BODY.BUSI_INFO.LOGIN_NO", "value": "E3ZZZZ063"},
                                           {"name": "ROOT.BODY.BUSI_INFO.COMMON_INFO.TRACE_ID",
                                            "value": "11*20250627143523**E3ZZZZ063*854463"},
                                           {"name": "ROOT.BODY.BUSI_INFO.COMMON_INFO.SRC_PORT", "value": ""},
                                           {"name": "ROOT.BODY.BUSI_INFO.COMMON_INFO.ROUND_AUDIT", "value": "N"}],
                        "sys.javaVersion": "1.6.0_45", "timepoint.LogFormatting": "2025-06-27 14:35:24.994",
                        "ulmp.node_lv1": "ah", "ulmp.node_lv2": "", "ulmp.no": "2025-06-27:550644", "source.file": "?",
                        "sys.clientAddress": "127.0.0.1", "sys.clientHostName": "localhost", "log.level": "INFO",
                        "source.method": "?", "timepoint.fmt.kafkaTopic": "inkey_a01(11)",
                        "log.message": "V.3.1.0.201009~!~2025-06-27 14:35:23,751~!~info~!~11*20250627143523**E3ZZZZ063*854463~!~2c6b99bb-2668-46e8-bce5-e961eaf6f469~!~~!~A101A~!~/bss-base-app/v1/crm/base/baseApp/workArea/qryWorkConfByLoginNo~!~~!~~!~E3ZZZZ063~!~82.221.145.202~!~10.243.141.71~!~11000~!~~!~~!~~!~~!~-2050919316~!~n466B175ABB796981712A82058B926575-1~!~-1235267275~!~~!~tangmaolong01~!~N~!~~!~~!~~!~~!~~!~~!~~!~统一业务受理平台~!~+20250627143523751+E3ZZZZ063+11*20250627143523**E3ZZZZ063*854463~!~0~!~a01~!~pin~!~{\"ROOT\":{\"BODY\":{\"BUSI_INFO\":{\"DESK_ID\":null,\"COMMON_INFO\":{\"SESSION_ID\":\"n466B175ABB796981712A82058B926575-1\",\"TRACE_ID\":\"11*20250627143523**E3ZZZZ063*854463\",\"CALL_ID\":\"2c6b99bb-2668-46e8-bce5-e961eaf6f469\",\"SER_NAME\":\"tangmaolong01\",\"ROUND_AUDIT\":\"N\",\"DEST_IP\":\"10.243.39.18\",\"DEST_PORT\":\"31000\",\"APP_NAME\":\"windows_IE11\",\"APP_ID\":\"\",\"FROM_SYS\":\"\",\"BACK2\":\"\",\"PROVINCE_GROUP\":\"10017\",\"SRC_PORT\":\"\"},\"OPR_INFO\":{\"OP_CODE\":\"\",\"LOGIN_NO\":\"E3ZZZZ063\",\"REGION_ID\":\"15\",\"CHANNEL_TYPE\":\"431\",\"LOGIN_PWD\":\"\",\"IP_ADDRESS\":\"82.221.145.202\",\"GROUP_ID\":\"115010975\",\"AUTHEN_CODE\":\"\",\"AUTHEN_NAME\":\"\",\"CONTACT_ID\":\"\",\"REL_OP_CODE\":\"\",\"CHANNEL_SOURCE\":\"\",\"RE_WAY\":\"\"},\"LOGIN_NO\":\"E3ZZZZ063\"}}}}",
                        "ulmp.logfmt": "002", "timepoint.fmt.kafkaOffset": "6037561537", "ulmp.inkey": "sk_026_01",
                        "ulmp.error_msg": "正则匹配不通过", "ulmp.log4j": "1", "source.class": "?",
                        "sys.ip": ["152.57.33.244", "10.243.141.71"],
                        "log.thread": "[ACTIVE] ExecuteThread: '165' for queue: 'weblogic.kernel.Default (self-tuning)'",
                        "ulmp.node_lv3": "sk_026", "ulmp.id": "dd1b5d6dcd00464e9d9d151e15537151", "sys.user": "echnop",
                        "old.ulmp.inkey": "a01", "fields": [
            {"name": "Log_ver", "remark": "日志版本", "type": "e2eLogType", "value": "V.3.1.0.201009"},
            {"name": "Timestamp", "remark": "时间戳", "type": "e2eTimestamp", "value": "2025-06-27 14:35:23,751"},
            {"name": "Log_level", "remark": "日志级别", "type": "logLevel", "value": "info"},
            {"name": "Trace_id", "remark": "统一流水", "type": "e2eTraceID",
             "value": "11*20250627143523**E3ZZZZ063*854463"},
            {"name": "Call_id", "remark": "调用ID", "type": "e2eCallID",
             "value": "2c6b99bb-2668-46e8-bce5-e961eaf6f469"},
            {"name": "Parent_call_id", "remark": "父调用ID", "type": "e2ePCallID", "value": ""},
            {"name": "Link_type", "remark": "环节类型", "type": "e2eLinkType", "value": "A101A"},
            {"name": "SpanName", "remark": "服务名", "type": "e2eSpanName",
             "value": "/bss-base-app/v1/crm/base/baseApp/workArea/qryWorkConfByLoginNo"},
            {"name": "return_code", "remark": "返回码", "type": "e2eResultCode", "value": ""},
            {"name": "Return_info", "remark": "返回信息", "type": "e2eResultInfo", "value": ""},
            {"name": "Login_no", "remark": "登陆工号", "type": "e2eUserID", "value": "E3ZZZZ063"},
            {"name": "Client_ip", "remark": "客户端IP", "type": "e2eClientIP", "value": "82.221.145.202"},
            {"name": "Host_ip", "remark": "服务端主机ip", "type": "e2eHostIP", "value": "10.243.141.71"},
            {"name": "Port", "remark": "服务端口", "type": "e2eHostPort", "value": "11000"},
            {"name": "Route_key", "remark": "路由类型", "type": "e2eObjectType", "value": ""},
            {"name": "Route_value", "remark": "路由值", "type": "e2eObjectValue", "value": ""},
            {"name": "Current_thread", "remark": "线程名", "type": "string", "value": ""},
            {"name": "Class_name", "remark": "类名", "type": "string", "value": ""},
            {"name": "tid", "remark": "访问id", "type": "string", "value": "-2050919316"},
            {"name": "sid", "remark": "会话id", "type": "string", "value": "n466B175ABB796981712A82058B926575-1"},
            {"name": "rid", "remark": "流程id", "type": "string", "value": "-1235267275"},
            {"name": "sql", "remark": "sql", "type": "string", "value": ""},
            {"name": "opr_ser_name_4A", "remark": "4A主账号", "type": "string", "value": "tangmaolong01"},
            {"name": "opr_round_audit_4A", "remark": "是否绕行", "type": "string", "value": "N"},
            {"name": "sessionId_4A", "remark": "金库会话ID", "type": "string", "value": ""},
            {"name": "startDate_4A", "remark": "有效开始时间", "type": "string", "value": ""},
            {"name": "endDate_4A", "remark": "有效结束时间", "type": "string", "value": ""},
            {"name": "Authen_type_OPcode", "remark": "用户授权凭证类型", "type": "string", "value": ""},
            {"name": "Operate_object", "remark": "操作对象", "type": "string", "value": ""},
            {"name": "Op_code", "remark": "模块代码", "type": "string", "value": ""},
            {"name": "operName", "remark": "操作名称", "type": "string", "value": ""},
            {"name": "resName", "remark": "资源名称", "type": "string", "value": "统一业务受理平台"},
            {"name": "operContent", "remark": "操作内容", "type": "string",
             "value": "+20250627143523751+E3ZZZZ063+11*20250627143523**E3ZZZZ063*854463"},
            {"name": "operResult", "remark": "操作结果", "type": "string", "value": "0"},
            {"name": "accessKey", "remark": "接入标识", "type": "string", "value": "a01"},
            {"name": "jcfParam", "remark": "出入参标识", "type": "e2eCallType", "value": "pin"},
            {"name": "in_param", "remark": "出入参字段", "type": "e2eInOutParams",
             "value": "{\"ROOT\":{\"BODY\":{\"BUSI_INFO\":{\"DESK_ID\":null,\"COMMON_INFO\":{\"SESSION_ID\":\"n466B175ABB796981712A82058B926575-1\",\"TRACE_ID\":\"11*20250627143523**E3ZZZZ063*854463\",\"CALL_ID\":\"2c6b99bb-2668-46e8-bce5-e961eaf6f469\",\"SER_NAME\":\"tangmaolong01\",\"ROUND_AUDIT\":\"N\",\"DEST_IP\":\"10.243.39.18\",\"DEST_PORT\":\"31000\",\"APP_NAME\":\"windows_IE11\",\"APP_ID\":\"\",\"FROM_SYS\":\"\",\"BACK2\":\"\",\"PROVINCE_GROUP\":\"10017\",\"SRC_PORT\":\"\"},\"OPR_INFO\":{\"OP_CODE\":\"\",\"LOGIN_NO\":\"E3ZZZZ063\",\"REGION_ID\":\"15\",\"CHANNEL_TYPE\":\"431\",\"LOGIN_PWD\":\"\",\"IP_ADDRESS\":\"82.221.145.202\",\"GROUP_ID\":\"115010975\",\"AUTHEN_CODE\":\"\",\"AUTHEN_NAME\":\"\",\"CONTACT_ID\":\"\",\"REL_OP_CODE\":\"\",\"CHANNEL_SOURCE\":\"\",\"RE_WAY\":\"\"},\"LOGIN_NO\":\"E3ZZZZ063\"}}}}"}],
                        "sys.dir": "/echnop/domains/echddomain", "log.timeMillis": 1751006123751, "log.length": 1164,
                        "log.loggerName": "CsdFlumeLogger", "timepoint.fmt.kafkaTimestamp": "2025-06-27 14:35:23.988"}

    # 初始化生产者
    producer = SecureKafkaProducer(kafka_config)
    total_sent = 0

    try:
        for i in range(args.iterations):
            start_time = time.time()

            # 更新消息模板中的序列号
            message_template['metadata']['sequence'] = i + 1

            # 发送批次
            success = producer.send_batch(message_template, args.batch_size)
            total_sent += success

            logger.info(f"Batch {i + 1}/{args.iterations} completed. "
                        f"Sent {success}/{args.batch_size} messages. "
                        f"Total sent: {total_sent}")

            # 控制发送频率
            elapsed = time.time() - start_time
            if elapsed < args.interval:
                time.sleep(args.interval - elapsed)

    except KeyboardInterrupt:
        logger.info("Producer interrupted by user")
    except Exception as e:
        logger.error(f"Fatal error: {str(e)}")
    finally:
        producer.producer.close()
        logger.info(f"Producer closed. Total messages sent: {total_sent}")


if __name__ == '__main__':
    main()
